package com.jiangsong.batch;


import com.jiangsong.entity.Product;
import org.springframework.batch.core.Job;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepScope;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
import org.springframework.batch.item.ItemWriter;
import org.springframework.batch.item.database.JdbcPagingItemReader;
import org.springframework.batch.item.database.JpaItemWriter;
import org.springframework.batch.item.database.Order;
import org.springframework.batch.item.database.builder.JdbcPagingItemReaderBuilder;
import org.springframework.batch.item.database.support.MySqlPagingQueryProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.SimpleAsyncTaskExecutor;
import org.springframework.jdbc.core.BeanPropertyRowMapper;

import javax.persistence.EntityManagerFactory;
import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Executors;

@Configuration
public class BatchConfig {
    private final JobBuilderFactory jobBuilderFactory;
    private final StepBuilderFactory stepBuilderFactory;

    @Autowired
    private DataSource dataSource; // 假设已经配置了 DataSource


    public BatchConfig(JobBuilderFactory jobBuilderFactory, StepBuilderFactory stepBuilderFactory) {
        this.jobBuilderFactory = jobBuilderFactory;
        this.stepBuilderFactory = stepBuilderFactory;
    }

    @Bean
    @StepScope
    public JdbcPagingItemReader<Product> jdbcPagingItemReader(
            @Value("#{stepExecutionContext['jobZoneNo']}") String jobZoneNo
    ) {
        Map<String, Order> sortKeys = new HashMap<>();
        sortKeys.put("acc_no", Order.DESCENDING); // 排序字段和排序顺序

        MySqlPagingQueryProvider queryProvider = new MySqlPagingQueryProvider();
        queryProvider.setSelectClause("SELECT distinct acc_no");
        queryProvider.setFromClause("FROM product");
        queryProvider.setWhereClause("WHERE valid_flag = :validFlag and job_zone_no = :jobZoneNo"); // 自定义 WHERE 子句
        queryProvider.setSortKeys(sortKeys);

        Map<String, Object> paramterMap = new HashMap<>();
        paramterMap.put("jobZoneNo", jobZoneNo);
        paramterMap.put("validFlag", "1");
        System.out.println("执行reader，jobNo:" + jobZoneNo);
        return new JdbcPagingItemReaderBuilder<Product>()
                .name("jdbcPagingItemReader")
                .dataSource(dataSource)
//                .fetchSize(10) // 每次读取的记录数
//                .rowMapper(new ProductRowMapper())
                .rowMapper(new BeanPropertyRowMapper<>(Product.class))
                .queryProvider(queryProvider)
                .parameterValues(paramterMap) // 自定义参数
                .pageSize(10) // 分页大小
                .build();
    }

    @Bean
    @StepScope
    public ProductItemProcessor itemProcessor(@Value("#{stepExecutionContext['jobZoneNo']}") String jobZoneNo
    ) {
        return new ProductItemProcessor();
    }

    @Bean
    public JpaItemWriter<Product> writer(EntityManagerFactory entityManagerFactory) {
        JpaItemWriter<Product> writer = new JpaItemWriter<>();
        writer.setEntityManagerFactory(entityManagerFactory);
        return writer;
    }

    @Bean
    public ItemWriter<Product> itemWriter() {
        return items -> {
//            for (MyObject item : items) {
//                System.out.println(item);
//            }
        };
    }

    @Bean
    public Step step1(JpaItemWriter<Product> writer) {
        return stepBuilderFactory.get("step1")
                .<Product, Product>chunk(5)
                .reader(jdbcPagingItemReader(null))
                .processor(itemProcessor(null))
                .writer(itemWriter())
//                .listener()
//                .listener(productChunkListener)
                .build();
    }

    @Bean
    public Job job(Step masterStep) {
        return jobBuilderFactory.get("job")
                .start(masterStep)
                .listener(new ProductJobListener())
                .build();
    }

    @Bean
    public ProductPartitioner productPartitioner() {
        return new ProductPartitioner();
    }

    @Bean
    public PartitionHandler partitionHandler(Step step1) {
        TaskExecutorPartitionHandler handler = new TaskExecutorPartitionHandler();

        handler.setGridSize(10);
        handler.setTaskExecutor(new SimpleAsyncTaskExecutor());
        handler.setStep(step1);

        return handler;
    }

    @Bean
    public Step masterStep(Step step1) {
        return stepBuilderFactory.get("masterStep")
                .partitioner(step1.getName(), productPartitioner())
                .partitionHandler(partitionHandler(step1))
                .build();
    }
}